import numpy as np
import pandas as pd
from sklearn.model_selection import train_test_split
import colorsys

# 样本数量
num_samples = 20000

# 生成随机的 Arousal 和 Valence 值，范围是 [-1, 1]
np.random.seed(42)
arousal_values = np.random.uniform(-1, 1, num_samples)
valence_values = np.random.uniform(-1, 1, num_samples)

# 生成模拟的 HSB 值（根据 AV 值进行模拟）
def generate_hsb(arousal, valence):
    hue = (2 / 3 - (valence + 1) / 3)
    saturation = (arousal + 1) / 2
    brightness = 1
    return hue, saturation, brightness

hsb_values = np.array([generate_hsb(a, v) for a, v in zip(arousal_values, valence_values)])

# 创建数据集表格
data = {
    "Arousal": arousal_values,
    "Valence": valence_values,
    "Hue": hsb_values[:, 0],
    "Saturation": hsb_values[:, 1],
    "Brightness": hsb_values[:, 2]
}

df = pd.DataFrame(data)

# 预处理数据
X = df[['Arousal', 'Valence']].values
# 只选择 Saturation 作为输出
S = df[['Saturation']].values

# 划分训练集和测试集
X_train, X_test, S_train, S_test = train_test_split(X, S, test_size=0.2, random_state=42)

# 合并输入和输出数据
train_data = np.hstack((X_train, S_train))
test_data = np.hstack((X_test, S_test))

# 将数据保存为TXT文件，并在每行前面添加行号
def save_to_txt_with_lineno(filename, data):
    with open(filename, 'w') as f:
        for i, row in enumerate(data, start=1):
            line = ','.join(map(str, row))
            f.write(f"{i},{line}\n")

# 保存为TXT文件
save_to_txt_with_lineno('train of S_dataset.txt', train_data)
save_to_txt_with_lineno('test of S_dataset.txt', test_data)